g = {
    1: [2, 4, 7, 11],
    2: [5, 10, 16, 20],
    3: [8, 12, 17, 22],
}

import math


def SquareRootFloor(x):
    return int(math.sqrt(x))


MAX_Y = 10
MAX_K = 3
MAX_X = 3  # g() 接受的最大x值


def DP():
    dp = {}
    # 记录当子问题取最大值时，给gk分配了的x多大。
    solution = {}

    for y in range(0, MAX_Y + 1):
        dp[1, y] = g[1][SquareRootFloor(y)]
        solution[1, y] = SquareRootFloor(y)
    sum = 0
    for k in range(1, MAX_K + 1):
        sum += g[k][0]
        dp[k, 0] = sum
        solution[k, 0] = 0

    for k in range(2, MAX_K + 1):
        for y in range(1, MAX_Y + 1):
            print('k', k, 'y', y)
            val = -1
            sol = None
            for xk in range(0, SquareRootFloor(y) + 1):
                temp3 = dp[k - 1, y - xk * xk] + g[k][xk]
                if temp3 > val:
                    val = temp3
                    sol = xk

            dp[k, y] = val
            solution[k, y] = sol
            print('dp', val, 'sol', sol)

    max_val = dp[MAX_K, MAX_Y]
    return max_val


def TrackSolution(solution, dp, K, Y):
    ans = []
    y = Y
    for k in range(K, 0, -1):
        x = solution[k, y]
        y -= x * x
        ans.append(x)

    ans.reverse()
    print(dp[K, Y], ans)
    return ans


def PrintTable(dp):
    for k in range(1, MAX_K + 1):
        for y in range(0, MAX_Y + 1):
            print(f' {dp[k,y]}', end='')
        print()


DP()
